package com.huai.jt1078.utils;

import cn.hutool.core.util.ArrayUtil;
import cn.hutool.core.util.ObjectUtil;
import cn.hutool.core.util.StrUtil;
import com.huai.jt1078.codec.AudioCodec;
import com.huai.jt1078.codec.MP3Encoder;
import com.huai.jt1078.config.VideoConfig;
import com.huai.jt1078.entity.AudioTag;
import com.huai.jt1078.entity.Media;
import com.huai.jt1078.entity.MediaEncoding;
import com.huai.jt1078.entity.VisualAudio;
import com.huai.jt1078.enums.NaluType;
import com.huai.jt1078.handler.thread.HeartbeatThread;
import com.huai.jt1078.handler.thread.PushVideoThread;
import com.huai.jt1078.handler.thread.TakeThread;
import io.netty.buffer.ByteBuf;
import io.netty.buffer.ByteBufAllocator;
import io.netty.buffer.Unpooled;
import lombok.Getter;
import lombok.extern.slf4j.Slf4j;

import java.io.ByteArrayOutputStream;
import java.util.LinkedList;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;

/**
 * @author xingkong
 * @program jt1078
 * @description 视频工具类
 * @date 2021-09-06 20:49
 **/
@Slf4j

public class VideoUtil {

    /** 输出字节流 */
    private final String tag;

    private byte[] data = new byte[0];

    @Getter
    private byte[] sps = null;
    private byte[] pps = null;

    private boolean writeAvcSeqHeader = false;
    private byte[] flvHeader;
    private byte[] flvAudioHeader;
    private byte[] videoHeader;
    ByteArrayOutputStream videoFrame;
    private byte[] lastIFrame;
    private boolean videoHeaderSent = false;
    private final Object lock;
    protected LinkedList<byte[]> messages;
    @Getter
    private AudioCodec audioCodec;
    private final MP3Encoder mp3Encoder;
    private final TakeThread takeThread;
    private final PushVideoThread pushVideoThread;
    private final HeartbeatThread heartbeatThread;

    private long videoTimestamp = 0;
    private long audioTimestamp = 0;

    /** 这里取其默认值0 */
    private final int streamId = 0;

    /**
     * 第一个时间戳
     */
    private Long firstTimestamp = -1L;

    /**
    * 上一个时间时间戳
    */
    private long lastVideoFrameTimeOffset;
    private long lastAudioFrameTimeOffset;

    private final ScheduledExecutorService scheduledExecutorService = Executors.newSingleThreadScheduledExecutor();

    /** nalu最小长度 */
    private final static int MIN_NALU_LENGTH = 1;

    public VideoUtil(String tag,VideoConfig videoConfig) {
        this.tag = tag;
        this.lock = new Object();
        this.messages = new LinkedList<byte[]>();
        videoFrame = new ByteArrayOutputStream(2048 * 100);
        makeFlvHeader();
        makeFlvAudioHeader();
        mp3Encoder = new MP3Encoder();

        // 启动推送视频流
        pushVideoThread = new PushVideoThread(tag,videoConfig);
        pushVideoThread.setDaemon(true);
        pushVideoThread.start();

        // 启动接收视频流
        this.takeThread = new TakeThread(lock, tag, messages);
        takeThread.setDaemon(true);
        takeThread.start();
        log.info("开始推送视频流:{}",tag);
        log.info("flv播放地址:{}" , videoConfig.getFlvUrl().replace("%s",tag));
       //启动心跳线程
        heartbeatThread = new HeartbeatThread(tag);
        heartbeatThread.setDaemon(true);
        if (videoConfig.isHeartbeat()){
            heartbeatThread.start();
        }
    }


    /**
     * 视频推送
     * @param visualAudio 视频推送实体
     */
    public void publishVideo(VisualAudio visualAudio){
        // 写入数据
        write(visualAudio.getData());
        // 如果当前尚未开始 则保存第一个时间戳
        if (SessionManager.DEFAULT_TIMESTAMP.equals(firstTimestamp)){
            firstTimestamp = visualAudio.getTimeStamp();
        }
        // nalu每个数据开头都是00000001
        String h264Header = "00000001";

        String dataHex = ByteUtil.bytes2HexStr(data);
        String[] dataHexList = dataHex.split(h264Header);
        // 把最后一段数据放入缓存区
        if (StrUtil.isNotBlank(dataHexList[dataHexList.length - 1])){
            this.data = ByteUtil.hexStr2Bytes(dataHexList[dataHexList.length - 1]);
        }else{
            this.data = new byte[0];
        }
        // 遍历数据
        for (int i = 0;i< dataHexList.length - 1;i++){
            byte[] bytes = ByteUtil.hexStr2Bytes(dataHexList[i]);
            if (bytes.length < 1){
                continue;
            }
            byte[] flvTag = getFlvTag(bytes, (visualAudio.getTimeStamp() - firstTimestamp));
            if (ObjectUtil.isNull(flvTag) || flvTag.length < 1){
                continue;
            }
            onVideoData(flvTag,visualAudio.getTimeStamp());

        }
    }


    /**
     * 音频推送
     * @param visualAudio 解析好的实体
     * @author xingkong
     * @date 2021-09-08 18:21
     */
    public void publishAudio(VisualAudio visualAudio) {

        if (ObjectUtil.isNull(audioCodec)){
            audioCodec = AudioCodec.getCodec(visualAudio.getLoadType());
            log.info("音频格式: {}", MediaEncoding.getEncoding(Media.Type.Audio, visualAudio.getLoadType()));
        }
        byte[] data = audioCodec.toPCM(visualAudio.getData());


        // 如果当前还没有推送视频头 则不播放声音
        if (!videoHeaderSent){
            // 判断当前tag是否是音频通道 如果是 则推送音频头 并且打开标识
            if (SessionManager.instance.getAudioMap().contains(tag)){
                enqueue(ByteUtil.make(flvAudioHeader));
                videoHeaderSent = true;
                SessionManager.instance.getAudioMap().remove(tag);
                SessionManager.instance.getWebsocketUtilMap().put(tag,new WebsocketUtil(tag));
            }
            return;
        }
        // 把音频转换为map3
        byte[] mp3Data = mp3Encoder.encode(data);
        // 如果音频没有数据 则直接跳过
        if (ObjectUtil.isNull(mp3Data) || mp3Data.length < 1){
            return;
        }
        // 构建一个音频tag
        AudioTag audioTag = new AudioTag(0, mp3Data.length + 1, AudioTag.MP3, (byte) 0, (byte)1, (byte) 0, mp3Data);
        byte[] frameData = null;
        try {
            ByteBuf audioBuf = ByteUtil.encode(audioTag);
            frameData = ByteUtil.readReadableBytes(audioBuf);
        } catch (Exception e) {
            e.printStackTrace();
        }
        // 获取时间戳
        long timeOffset = visualAudio.getTimeStamp();
        // 如果尚未播放过音频则开始播放音频
        if (lastAudioFrameTimeOffset == 0) {
            lastAudioFrameTimeOffset = timeOffset;
        }
        // 如果没有数据则直接返回
        if (ObjectUtil.isNull(frameData) || frameData.length < 1){
            return;
        }
        resetTimestamp(frameData, (int) audioTimestamp);
        audioTimestamp += (int)(timeOffset - lastAudioFrameTimeOffset);
        lastAudioFrameTimeOffset = timeOffset;

        enqueue(ByteUtil.make(frameData));
    }


    ////////////////////// 以下为私有方法   /////////////////////////////


    /**
    * 设置flv的header
    * @author xingkong
    * @date 2021/9/8
    */
    private void makeFlvHeader(){
        ByteBuf buffer = ByteBufAllocator.DEFAULT.buffer(13);
        // 设置signature
        buffer.writeByte((byte)'F');
        buffer.writeByte((byte)'L');
        buffer.writeByte((byte)'V');
        // 设置版本
        buffer.writeByte(0x01);
        // 设置flag 0x01是有视频 0x04是有声音 这里写有视频 去 | TODO 后续可以测试 如果是0x05会怎么样
        buffer.writeByte(0x01);
        // headerSize 默认为9
        buffer.writeInt(0x09);
        // 以上才是真正的header 下面是为了方便后续拼接报文 所以先设置了一个previous Tag size

        // 因为第一个previous Tag size是0  原本是 【previous Tag size + Tag】的组合 这里修改为 0 + 【Tag + previous Tag size】的组合
        buffer.writeInt(0x00);

        flvHeader = ByteUtil.readReadableBytes(buffer);
    }

    /**
     * 如果仅有音频 则设置flv头为这个
     * @author xingkong
     * @date 2021-09-17 14:14
     */
    private void makeFlvAudioHeader(){
        ByteBuf buffer = ByteBufAllocator.DEFAULT.buffer(13);
        // 设置signature
        buffer.writeByte((byte)'F');
        buffer.writeByte((byte)'L');
        buffer.writeByte((byte)'V');
        // 设置版本
        buffer.writeByte(0x01);
        // 设置flag 0x01是有视频 0x04是有声音 这里写有音频
        buffer.writeByte(0x04);
        // headerSize 默认为9
        buffer.writeInt(0x09);
        // 以上才是真正的header 下面是为了方便后续拼接报文 所以先设置了一个previous Tag size

        // 因为第一个previous Tag size是0  原本是 【previous Tag size + Tag】的组合 这里修改为 0 + 【Tag + previous Tag size】的组合
        buffer.writeInt(0x00);

        flvAudioHeader = ByteUtil.readReadableBytes(buffer);
    }

    /**
    * 压入数据到最后
    * @param  nalu 待压入数据
    * @author xingkong
    * @date 2021/9/8
    */
    private void enqueue(byte[] nalu){
        if (ObjectUtil.isNull(nalu)){
            return;
        }
        // 抢锁并且压入数据到最后
        synchronized (lock){
            messages.addLast(nalu);
            lock.notify();
        }
    }

    /**
    * 发送数据
    * @param data flv视频格式数据
     * @param timeOffset 时间戳
    * @author xingkong
    * @date 2021/9/8
    */
    private void onVideoData(byte[] data,long timeOffset) {
        // 如果是第一次推送 则把当前时间戳赋值给上一帧时间
        if (0 == lastVideoFrameTimeOffset){
            lastVideoFrameTimeOffset = timeOffset;
        }
        // 如果之前没有发送过videoHeader 并且 header已经准备好了 那就发送header


        if (!videoHeaderSent && writeAvcSeqHeader){
            enqueue(ByteUtil.make(flvHeader));
            enqueue(ByteUtil.make(videoHeader));

            // 记录发送视频的高度和宽度
            log.info("视频的高度：{} 视频的宽度：{}",ByteUtil.getHeight(sps),ByteUtil.getWidth(sps));

            // 直接发送第一帧
            if (ObjectUtil.isNotNull(lastIFrame)){
                resetTimestamp(lastIFrame,videoTimestamp);
                enqueue(ByteUtil.make(lastIFrame));
            }
            videoHeaderSent = true;
        }

        resetTimestamp(data,videoTimestamp);
        videoTimestamp += (int)(timeOffset - lastVideoFrameTimeOffset);
        lastVideoFrameTimeOffset = timeOffset;
        // 压入数据
        enqueue(ByteUtil.make(data));

    }

    /**
    *  充值时间
    * @param  packet nalu数据
     * @param timestamp 时间戳
    * @author xingkong
    * @date 2021/9/8
    */
    private void resetTimestamp(byte[] packet, long timestamp) {
        // 只对视频类的TAG进行修改
        if (packet[0] != 9 && packet[0] != 8) {
            return;
        }
        packet[4] = (byte)((timestamp >> 16) & 0xff);
        packet[5] = (byte)((timestamp >>  8) & 0xff);
        packet[6] = (byte)((timestamp) & 0xff);
        packet[7] = (byte)((timestamp >> 24) & 0xff);
    }

    /**
     * 获取 flv视频的tag
     * @param nalu h264 最小单元
     * @param timestamp 时间差
     * @return flv的tag
     */
    private byte[] getFlvTag(byte[] nalu,Long timestamp){

        // 如果小于最小nalu长度 或者时间戳丢失 则不需要该数据
        if (nalu.length <= MIN_NALU_LENGTH || ObjectUtil.isNull(timestamp)){
            return new byte[0];
        }

        // NAL头 0位是F 初始为0 1~2位保存这NRI 标识重要等级（暂无实用） 3~7位保存NAL类型
        // 0 未指定
        // 1 非IDR的片
        // 2 片数据A分区
        // 3 片数据B分区
        // 4 片数据C分区
        // 5 IDR图像的片
        // 6 补充增强信息单元（SEI）
        // 7 序列参数集（SPS）
        // 8 图像参数集（PPS）
        // 9 分界符
        // 10 序列结束
        // 11 码流结束
        int naluType = nalu[0] & 0x1f;

        // 这里仅处理类型为1 2 3 4 5 7 8
        if (!NaluType.naluTypes().contains(naluType)){
            return null;
        }

        // 如果当前没有设置SPS且 当前nalu类型为0x07
        if (null == sps &&  NaluType.SPS.getCode().equals(naluType)){
            this.sps = nalu;
        }

        // 如果当前没有设置pps且 当前nalu类型为0x08
        if (null == pps &&  NaluType.PPS.getCode().equals(naluType)){
            this.pps = nalu;
        }

        // 如果sps pps都有了 且还没有写入头 则开始写入头信息
        if (ObjectUtil.isNotNull(pps) && ObjectUtil.isNotNull(sps) && (!writeAvcSeqHeader)){
            writeH264Header(timestamp);
            // 写入头信息之后 修改写入头为true
            writeAvcSeqHeader = true;
        }

        // 如果还没有写avs视频头 则跳过 一个视频可能有多个SPS和PPS 我们这里可以认为只有一个SPS和PPS 这里不处理SPS和PPS
        if (!writeAvcSeqHeader){
            return null;
        }

        videoFrame.reset();
        // 开始写入 h264 frame数据
        writeH264Frame(nalu,timestamp);

        // 如果frame数据为空 则直接返回
        if (videoFrame.size() < 1){
            return null;
        }

        // 如果当前NAL单元为I祯，则缓存该帧数据为上一帧
        if (NaluType.IDR.getCode() == naluType){
            this.lastIFrame = videoFrame.toByteArray();
        }

        return videoFrame.toByteArray();
    }



    /**
     * 写入数据
     * @param h264 h264数据
     */
    private void write(byte[] h264){
        this.data = ArrayUtil.addAll(this.data,h264);
    }

    /**
    * 写入h264 frame数据
    * @param nalu h264基本流组成单位
    * @author xingkong
    * @date 2021/9/7
    */
    private void writeH264Frame(byte[] nalu, Long timestamp) {
//        System.out.println("时间戳:"+timestamp);
        Integer naluType = nalu[0] & 0x1f;
        // 这里不处理SPS和PPS
        if (NaluType.PPS.getCode().equals(naluType) || NaluType.SPS.getCode().equals(naluType)){
            return;
        }
        // 创建一个bytebuf来保存frame
        ByteBuf buffer = Unpooled.buffer();

        /////////////////////////////  开始设置tag header /////////////////////////////
        // 0x08是音频 0x09是视频 这里写视频 【1】
        buffer.writeByte(0x09);
        // 前一个tag的长度 Flv的header由Signature[1] + Version [1] + Flags[3] + HeaderSize[4]
        // 本字段为数据长度 = Flv的header的长度 + nalu长度  【3】
        int dataSize = 1 + 1 + 3 + 4 + nalu.length;
        buffer.writeBytes(ByteUtil.getThreeBytes(dataSize));
        // 设置时间戳 【4】
        buffer.writeBytes(ByteUtil.getThreeBytes(timestamp.intValue()));
        buffer.writeByte((byte)(timestamp >> 24));
        // 设施streamId 【3】
        buffer.writeBytes(ByteUtil.getThreeBytes(0x00));

        ////////////////////////////// 开始设置Tag data ////////////////////////////////
        // 帧类型 + 编码ID  https://juejin.cn/post/6844903598476754951 【1】
        if (NaluType.IDR.getCode().equals(naluType)) {
            // 如果当前是关键帧  则帧类型为1（0001） + 编码类型是AVC = 7（0111）
            buffer.writeByte(0x17);
        } else {
            // 如果不是关键帧 则帧类型为2 （0010） + 编码类型是AVC = 7（0111）
            buffer.writeByte(0x27);
        }
        // AVC VIDEO PACKET = AVC packet类型 + CTS + 数据
        // 1：AVC NALU单元 【1】
        buffer.writeByte(0x01);
        // packet类型是0，则为0 3个字节 【3】
        buffer.writeBytes(ByteUtil.getThreeBytes(0x00));
        // 数据长度
        buffer.writeInt(nalu.length);
        // 数据信息
        buffer.writeBytes(nalu);
        // tag header的长度是11  dataSize是真实数据的大小
        buffer.writeInt(11 + dataSize);
        byte[] bytes = ByteUtil.readReadableBytes(buffer);
        videoFrame.write(bytes,0,bytes.length);
    }

    /**
    *
     *          DataSize=5 +                  // AVC video tag header (FrameType + CodecID | .. CompositionTime)
     *          6 +                   // AVCDecoderConfigurationRecord
     *          SPSCount*2 +               // 每个SPS长度2字节
     *          各个 SPSDataLength + // 所有SPS数据长度和
     *          1 +                                     // PPS个数
     *          PPSCount*2 +                // 每个PPS长度2字节
     *          各个 PPSDataLength;   // 所有PPS数据长度和
    * @param
    * @return
    * @author xingkong
    * @create 2021/9/7
    */
    private void writeH264Header(Long timestamp) {
        int dataSize = 5 + 6 + 2 + sps.length + 1 + 2 + pps.length ;
        // 创建中间变量保存数据
        ByteBuf buffer = ByteBufAllocator.DEFAULT.buffer(dataSize + 32);
        // TagType 【1】
        buffer.writeByte(0x09);
        // DataSize 【3】
        buffer.writeBytes(ByteUtil.getThreeBytes(dataSize));
        //  TimeStamp | TimeStampExten 【4】
        buffer.writeBytes(ByteUtil.getThreeBytes(timestamp.intValue()));
        buffer.writeByte((byte)(timestamp >> 24));
        // StreamID 【3】
        buffer.writeBytes(ByteUtil.getThreeBytes(streamId));
        // 以上11个字节


        // FrameType | CodecID 这里写默认值0x17 【1】
        buffer.writeByte(0x17);
        // AVCPacketTyp 这里写默认值为0x00 【1】
        buffer.writeByte(0x00);
        // CompositionTime 这里写默认值0x00 【3】
        buffer.writeBytes(ByteUtil.getThreeBytes(0x00));
        // 以上5个字节


        // ConfigurationVersion 这里写默认版本0x01 【1】
        buffer.writeByte(0x01);
        // AVC Profile SPS[1]  【1】
        buffer.writeByte(sps[1]);
        // profile_compatibility  SPS[2] 【1】
        buffer.writeByte(sps[2]);
        //  AVC Level  SPS[3] 【1】
        buffer.writeByte(sps[3]);
        // lengthSizeMinusOne 这里写默认值0xff 【1】
        buffer.writeByte(0xff);
        // numOfSequenceParameterSets, reserved 3bits | SPS count, commonly be 1 这里写默认值0xe1 【1】
        buffer.writeByte(0xe1);
        // 以上6个字节

        // SPS0 Length N0 Byte 【2】
        buffer.writeShort(sps.length);
        // N0 Byte  SPS0 Data 循环存放最多31个SPS 【sps.length】
        buffer.writeBytes(sps);
        // 以上 2+ sps长度 个字节

        // PPS count 我们只传一个 【1】
        buffer.writeByte(0x01);
        // 以上一个字节

        // PPS0 length 【2】
        buffer.writeShort(pps.length);
        // PPS0 data
        buffer.writeBytes(pps);
        // 以上 2 + pps 个字节

        // 这个tag的 PreviousTagSize = 11 + DataSize。 11 是Video tag （TagType到StreamID）
        buffer.writeInt(11 + dataSize);

        videoHeader = ByteUtil.readReadableBytes(buffer);
    }

    public void close(){
        this.heartbeatThread.interrupt();
        this.pushVideoThread.interrupt();
        this.takeThread.interrupt();
    }


}